#!/bin/bash

# Spring Boot分库分表性能测试脚本

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 配置
BASE_URL="http://localhost:9999"
CONCURRENT_USERS=100
TOTAL_REQUESTS=1000
DURATION_SECONDS=60

# 检查依赖
command -v curl >/dev/null 2>&1 || { echo -e "${RED}curl未安装，请先安装curl${NC}"; exit 1; }
command -v jq >/dev/null 2>&1 || { echo -e "${YELLOW}jq未安装，部分功能可能受限${NC}"; }

# 创建测试数据
create_test_data() {
    echo -e "${GREEN}创建测试数据...${NC}"
    
    # 创建100个测试用户
    for i in {1..100}; do
        curl -s -X POST "$BASE_URL/api/users" \
            -H "Content-Type: application/json" \
            -d "{\"username\":\"load_test_user_$i\",\"email\":\"user$i@test.com\",\"phone\":\"13800$(printf '%05d' $i)\",\"age\":$((20 + i % 30))}" > /dev/null
    done
    
    echo -e "${GREEN}测试数据创建完成${NC}"
}

# 单用户查询测试
single_user_test() {
    echo -e "${GREEN}执行单用户查询测试...${NC}"
    
    local start_time=$(date +%s%N)
    local response=$(curl -s -w "%{http_code}" "$BASE_URL/api/users/count")
    local end_time=$(date +%s%N)
    
    local duration=$(( (end_time - start_time) / 1000000 ))
    local http_code=${response: -3}
    
    echo -e "单用户查询耗时: ${duration}ms, 状态码: $http_code"
}

# 并发插入测试
concurrent_insert_test() {
    echo -e "${GREEN}执行并发插入测试...${NC}"
    
    local start_time=$(date +%s)
    
    # 使用GNU parallel进行并发测试（如果可用）
    if command -v parallel >/dev/null 2>&1; then
        seq 1 $TOTAL_REQUESTS | parallel -j $CONCURRENT_USERS "
            curl -s -X POST $BASE_URL/api/users \
                -H 'Content-Type: application/json' \
                -d '{\"username\":\"parallel_user_{}\",\"email\":\"user{}@test.com\",\"phone\":\"13900$(printf '%05d' {})\",\"age\":$((20 + {} % 30))}' > /dev/null
        "
    else
        # 使用简单的后台进程进行并发测试
        for i in $(seq 1 $TOTAL_REQUESTS); do
            curl -s -X POST "$BASE_URL/api/users" \
                -H "Content-Type: application/json" \
                -d "{\"username\":\"batch_user_$i\",\"email\":\"batch$i@test.com\",\"phone\":\"13700$(printf '%05d' $i)\",\"age\":$((20 + i % 30))}" > /dev/null &
            
            if [ $((i % CONCURRENT_USERS)) -eq 0 ]; then
                wait
            fi
        done
        wait
    fi
    
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))
    local rps=$((TOTAL_REQUESTS / duration))
    
    echo -e "并发插入测试完成: ${TOTAL_REQUESTS}次请求, ${duration}s, ${rps} RPS"
}

# 查询性能测试
query_performance_test() {
    echo -e "${GREEN}执行查询性能测试...${NC}"
    
    local total_requests=0
    local total_time=0
    
    for i in {1..10}; do
        local start_time=$(date +%s%N)
        curl -s "$BASE_URL/api/users" > /dev/null
        local end_time=$(date +%s%N)
        
        local duration=$(( (end_time - start_time) / 1000000 ))
        total_time=$((total_time + duration))
        total_requests=$((total_requests + 1))
    done
    
    local avg_time=$((total_time / total_requests))
    echo -e "查询平均响应时间: ${avg_time}ms"
}

# 分片分布检查
sharding_distribution_check() {
    echo -e "${GREEN}检查分片数据分布...${NC}"
    
    local count=$(curl -s "$BASE_URL/api/users/count")
    echo -e "总用户数: $count"
    
    # 检查每个分片的用户分布（需要数据库连接）
    if command -v mysql >/dev/null 2>&1; then
        echo -e "${YELLOW}数据库分片分布:${NC}"
        mysql -hlocalhost -P503 -uroot -p503922 -e "
        SELECT 'sharding_db_0.t_user_0' as table_name, COUNT(*) as count FROM sharding_db_0.t_user_0
        UNION ALL SELECT 'sharding_db_0.t_user_1', COUNT(*) FROM sharding_db_0.t_user_1
        UNION ALL SELECT 'sharding_db_0.t_user_2', COUNT(*) FROM sharding_db_0.t_user_2
        UNION ALL SELECT 'sharding_db_0.t_user_3', COUNT(*) FROM sharding_db_0.t_user_3
        UNION ALL SELECT 'sharding_db_1.t_user_0', COUNT(*) FROM sharding_db_1.t_user_0
        UNION ALL SELECT 'sharding_db_1.t_user_1', COUNT(*) FROM sharding_db_1.t_user_1
        UNION ALL SELECT 'sharding_db_1.t_user_2', COUNT(*) FROM sharding_db_1.t_user_2
        UNION ALL SELECT 'sharding_db_1.t_user_3', COUNT(*) FROM sharding_db_1.t_user_3;
        " 2>/dev/null || echo "数据库连接失败，跳过分布检查"
    fi
}

# 主测试流程
main() {
    echo -e "${GREEN}Spring Boot分库分表性能测试开始${NC}"
    echo -e "测试配置: 并发用户=${CONCURRENT_USERS}, 总请求=${TOTAL_REQUESTS}, 持续时间=${DURATION_SECONDS}s"
    
    # 检查服务是否可用
    if ! curl -s "$BASE_URL/api/users/count" > /dev/null; then
        echo -e "${RED}错误: 服务不可用，请先启动应用${NC}"
        exit 1
    fi
    
    # 执行测试
    create_test_data
    single_user_test
    concurrent_insert_test
    query_performance_test
    sharding_distribution_check
    
    echo -e "${GREEN}性能测试完成${NC}"
}

# 执行主函数
main "$@"